package com.ai.main.controller;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

/**
 * 功能描述
 *
 * @author: jichao.huang
 * @date: 2025年06月05日 9:32
 */
@RequiredArgsConstructor
@RequestMapping("/api/vectorStore")
@RestController
public class VectorStoreController {
    private final ChatClient chatClient;

    private final VectorStore vectorStore;

    @GetMapping(value = "/load")
    public String execute() {
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
        vectorStore.add(documents);

        return "ok";
    }

    @GetMapping(value = "/search")
    public List<Document> search(String primaryKey) {
        List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder().query(primaryKey).topK(5).build());

        return results;
    }
}
